#include "DaviesMaterial.cuh"

using namespace OBR;

static __forceinline__ __device__ Spectrum calcFr(SamplePRD *prd, float *pdf,
                                                  DaviesMaterialData *data)
{
    InteractionData &si = prd->inter;
    float costhetai = max(0.0f, dot(si.wi.d, si.n));
    float costhetao = max(0.0f, dot(-si.wo.d, si.n));
    if (costhetai <= 0)
    {
        *pdf = 0.0f;
        return 0.0f;
    }
    if (costhetao <= 0)
    {
        *pdf = 0.0f;
        return 0.0f;
    }
    float cos2thetai = costhetai * costhetai;
    float cos2thetao = costhetao * costhetao;
    float sin2thetai = 1.0f - cos2thetai;
    float sin2thetao = 1.0f - cos2thetao;
    float sinthetai = sqrtf(sin2thetai);
    float sinthetao = sqrtf(sin2thetao);
    float cosdphi =
        clamp(dot(normalize(si.wi.d - si.n * costhetai), normalize(-si.wo.d - si.n * costhetao)),
              -1.0f, 1.0f);

    float rad = si.has_differentials ? sqrtf(min(si.dudx * si.dudx + si.dudy * si.dudy,
                                                 si.dvdx * si.dvdx + si.dvdy * si.dvdy))
                                     : 0;

    float rho = data->rho.sample(si.uv, rad, rad, MipmapSampleMode::SQUARE);
    float rms = data->rms.sample(si.uv, rad, rad, MipmapSampleMode::SQUARE) * 1e3;
    float a = data->a.sample(si.uv, rad, rad, MipmapSampleMode::SQUARE) * 1e3;
    float rmsSqr = rms * rms;
    float aSqr = a * a;

    // PDF
    *pdf = costhetai * M_1_PIf;

    // brdfs
    Spectrum mix(0.0); // mix between 2 brdf
    Spectrum f_smooth, f_rough;

    // precompute per wavelength data
    Spectrum invLambdaSqr(1.0f);
    for (size_t s = 0; s < Spectrum::nSamples; s++)
    {
        invLambdaSqr[s] /= spectrumWavelengths[s] * spectrumWavelengths[s];
        mix[s] = smoothstep(1.0f, 1.5f, rms / spectrumWavelengths[s]);
    }

    // speular area
    float HoN = dot(normalize(si.n), normalize(si.wi.d - si.wo.d));
    float specWeight = smoothstep(0.999, 1.000, HoN);

    //* from this paper: https://cdmd.cnki.com.cn/Article/CDMD-11914-2008173180.htm */
    Spectrum spec = specWeight * (132.0f * M_1_PIf / costhetai) *
                    (invLambdaSqr * (-16.0f * M_PI2f * cos2thetai * rmsSqr)).exp();
    Spectrum offSpec =
        invLambdaSqr * invLambdaSqr / (costhetai * costhetao) *
        (M_PI2f * M_PIf * aSqr * rmsSqr * powf(costhetai + costhetao, 4.0f)) *
        (invLambdaSqr *
         (-M_PI2f * aSqr * (sin2thetai + sin2thetao + 2.0f * sinthetai * sinthetao * cosdphi)))
            .exp();
    f_smooth = rho * (spec + offSpec) * costhetai;
    //** very rough surface */
    float exp_pow = -(aSqr / rmsSqr) * 0.25f *
                    (sin2thetai + sin2thetao + 2.0f * sinthetai * sinthetao * cosdphi) /
                    ((costhetai + costhetao) * (costhetai + costhetao));
    float v = 0.0625f / (costhetai * costhetao) * M_1_PIf * (aSqr / rmsSqr) * expf(exp_pow);
    f_rough = rho * Spectrum(v) * costhetai;

    return mix * f_rough + (Spectrum(1.0f) - mix) * f_smooth;
}

static __forceinline__ __device__ Spectrum rayGeneration(SamplePRD *prd, float *pdf,
                                                         DaviesMaterialData *data, float u0,
                                                         float u1)
{
    // sample hemisphere
    Onb onb(prd->inter.n);
    float3 hemisphere_dir = cosine_sample_hemisphere(u0, u1);
    prd->inter.wi = prd->inter.spawnRay(onb.local2world(hemisphere_dir));

    return calcFr(prd, pdf, data);
}

extern "C" __device__ Spectrum __direct_callable__davies_material(SamplePRD *prd, float *pdf,
                                                                  float u0, float u1,
                                                                  MaterialCallMode mode)
{
    // get sbt data
    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    DaviesMaterialData *data = reinterpret_cast<DaviesMaterialData *>(sbt_record->callable_payload);
    data = data + prd->inter.mat_instance_id;

    if (mode == MaterialCallMode::SAMPLE_FR)
    {
        return rayGeneration(prd, pdf, data, u0, u1);
    }
    else if (mode == MaterialCallMode::CALC_FR_PDF)
    {
        return calcFr(prd, pdf, data);
    }
    return Spectrum(0.0f);
}